﻿#include <opencv2/opencv.hpp>
#include <iostream>

// 把矩阵中心移到左上角（OpenCV 的 dft 默认格式）
void fftShift(cv::Mat& src)
{
    int cx = src.cols / 2;
    int cy = src.rows / 2;
    cv::Mat q0(src, cv::Rect(0, 0, cx, cy));       // Top-Left
    cv::Mat q1(src, cv::Rect(cx, 0, cx, cy));      // Top-Right
    cv::Mat q2(src, cv::Rect(0, cy, cx, cy));      // Bottom-Left
    cv::Mat q3(src, cv::Rect(cx, cy, cx, cy));     // Bottom-Right
    cv::Mat tmp;
    q0.copyTo(tmp); q3.copyTo(q0); tmp.copyTo(q3); // swap TL <-> BR
    q1.copyTo(tmp); q2.copyTo(q1); tmp.copyTo(q2); // swap TR <-> BL
}

// 生成水平运动模糊 PSF（单通道，居中）
cv::Mat createMotionPSF(int rows, int cols, int len)
{
    cv::Mat psf = cv::Mat::zeros(rows, cols, CV_32F);
    int cx = cols / 2, cy = rows / 2;
    for (int x = cx - len / 2; x <= cx + len / 2; ++x)
        psf.at<float>(cy, x) = 1.0f / len;
    return psf;
}

// 对单通道做 FFT
cv::Mat forwardFFT(const cv::Mat& src)
{
    cv::Mat planes[] = { cv::Mat_<float>(src), cv::Mat::zeros(src.size(), CV_32F) };
    cv::Mat complex;
    cv::merge(planes, 2, complex);
    cv::dft(complex, complex, cv::DFT_COMPLEX_OUTPUT);
    return complex;
}

// 单通道逆 FFT
cv::Mat inverseFFT(const cv::Mat& complex)
{
    cv::Mat inv;
    cv::idft(complex, inv, cv::DFT_REAL_OUTPUT | cv::DFT_SCALE);
    cv::Mat dst;
    inv.convertTo(dst, CV_8U);  
    return dst;
}

// Wiener 滤波核心：numerator ./ (denominator + K)
cv::Mat wienerFilter(const cv::Mat& blurFreq, const cv::Mat& psfFreq, float K)
{
    cv::Mat psfConj;
    cv::mulSpectrums(psfFreq, cv::Mat::ones(psfFreq.size(), psfFreq.type()), psfConj, true); 

    cv::Mat numerator, denominator;
    cv::mulSpectrums(blurFreq, psfConj, numerator, 0);
    cv::mulSpectrums(psfFreq, psfConj, denominator, 0);

    // 逐像素做 (a+ib)/(c+id) = (a+ib)(c-id)/(c²+d²+K)
    cv::Mat res(blurFreq.size(), blurFreq.type());
    for (int i = 0; i < blurFreq.rows; ++i)
    {
        for (int j = 0; j < blurFreq.cols; ++j)
        {
            cv::Vec2f n = numerator.at<cv::Vec2f>(i, j);
            cv::Vec2f d = denominator.at<cv::Vec2f>(i, j);
            float mag2 = d[0] * d[0] + d[1] * d[1];
            float denom = mag2 + K;
            res.at<cv::Vec2f>(i, j)[0] = (n[0] * d[0] + n[1] * d[1]) / denom;
            res.at<cv::Vec2f>(i, j)[1] = (n[1] * d[0] - n[0] * d[1]) / denom;
        }
    }
    return res;
}

// 对单通道做“模糊 + Wiener 恢复”
void processChannel(const cv::Mat& ch, const cv::Mat& psfFreq, float K,
    cv::Mat& blurredCh, cv::Mat& recoveredCh)
{
    cv::Mat chFreq = forwardFFT(ch);
    cv::Mat blurFreq;
    cv::mulSpectrums(chFreq, psfFreq, blurFreq, 0);
    blurredCh = inverseFFT(blurFreq);

    cv::Mat recoveredFreq = wienerFilter(blurFreq, psfFreq, K);
    recoveredCh = inverseFFT(recoveredFreq);
}

int main()
{
    std::string imgPath = "example_image.jpg";  
    cv::Mat bgr = cv::imread(imgPath);   
    if (bgr.empty())
    {
        std::cerr << "Cannot read image: " << imgPath << std::endl;
        return -1;
    }

    std::vector<cv::Mat> channels, blurredChannels, recoveredChannels;
    cv::split(bgr, channels);

    cv::Mat psf = createMotionPSF(bgr.rows, bgr.cols, 21);
    fftShift(psf);
    cv::Mat psfFreq = forwardFFT(psf);
    const float K = 0.5f;

    for (size_t i = 0; i < 3; ++i)
    {
        cv::Mat blurredCh, recoveredCh;
        processChannel(channels[i], psfFreq, K, blurredCh, recoveredCh);
        blurredChannels.push_back(blurredCh);
        recoveredChannels.push_back(recoveredCh);
    }

    cv::Mat blurred, recovered;
    cv::merge(blurredChannels, blurred);
    cv::merge(recoveredChannels, recovered);

    cv::imshow("Original", bgr);
    cv::imshow("Blurred", blurred);
    cv::imshow("Wiener recovered", recovered);
    cv::imwrite("blurred.jpg", blurred);
    cv::imwrite("recovered.jpg", recovered);
    cv::waitKey();
    return 0;
}